Skip to content

feat(training): GPU-resident fused step for non-TS single-net models#1843

Open
ooples wants to merge 17 commits into
masterfrom
feat/gpu-resident-nonts
Open

feat(training): GPU-resident fused step for non-TS single-net models#1843
ooples wants to merge 17 commits into
masterfrom
feat/gpu-resident-nonts

Conversation

@ooples

@ooples ooples commented Jul 10, 2026

Copy link
Copy Markdown
Owner

Closes #1844

Non-time-series GPU-residency sweep

Paired with Tensors PR ooples/AiDotNet.Tensors#763 (compiled backward + createGraph, DP-SGD helper, multi-slot persistent inputs).

Files converted (23 model files, 30+ training methods)

Shared helpers

  • src/Training/GpuResidentFusedStep.cs — fused-step entry point for classes outside NeuralNetworkBase / TimeSeriesModelBase, with extra-tensors support for raw trainable state.
  • src/Training/DpSgdStep.cs — local mirror of Tensors PR test: add ProgramSynthesis integration tests #763's DpSgdStep. Same API; when NuGet republishes, swap the using.

Neural Networks (Graph + PhysicsInformed)

  • GraphClassificationModel, LinkPredictionModel, FourierNeuralOperator

Finance forecasters

  • TFC, TOTEM, CSDI, MQCNN

Synthetic-data generators (generator + disc + VAE + reconstruction)

  • Generator side: PATEGAN, MedSynth, CausalGAN, OCTGAN, TableGAN, CTGAN, DPCTGAN, CopulaGAN, TVAE, TabSyn, MisGAN (data+mask), TimeGAN (P1+P2)
  • Discriminator side (WGAN-GP fused): CausalGAN, CTGAN, CopulaGAN, TableGAN — fused compiled plan on the full Wasserstein + λ·GP objective, activated by Tensors PR test: add ProgramSynthesis integration tests #763's compiled-backward + createGraph support.
  • DP-SGD critic: DPCTGAN, MedSynth — per-example forward+backward, clip-before-aggregate, Gaussian noise, average — routed through the shared DpSgdStep helper so the Abadi 2016 privacy proof's L2-sensitivity contract can't be silently reversed.

Extra trainable state (via 4A extra-tensors API)

  • GOGGLE — soft adjacency A
  • CLAP — learned _logTemperature scalar

Diffusion (batched-per-element foundation)

  • INoiseScheduler / NoiseSchedulerBase.AddNoiseBatched — batched-per-element forward diffusion (Ho 2020 canonical). Virtual with per-element scalar default for backward compat; concrete schedulers can override.
  • NoisePredictorBase.PredictNoiseBatched + DiffusionModelBase.PredictNoiseBatched — batched-per-element noise prediction virtuals with default slice-then-call-scalar.
  • DiffusionModelBase.Train — rank ≥ 2 inputs go through the batched path (industry-standard); rank-1 stays on scalar for backward compat.

AiDotNet issue #1844 — FIXED HERE

WGAN-GP gradient penalty across every SyntheticData GAN was broken (createGraph=false in inner ComputeGradients meant the penalty gradient never reached disc weights, silently degrading WGAN-GP to plain WGAN). Fixed in all 5 files (CTGAN, DPCTGAN, CopulaGAN, CausalGAN, TableGAN) with createGraph: true.

Test plan

  • Builds green net471 / net8.0 / net10.0 on every commit
  • Fallback contract preserved — every fused-plan wire-up falls through to the (now-corrected) eager tape+optimizer path when the fused plan can't engage
  • CI green
  • Once Tensors PR test: add ProgramSynthesis integration tests #763 lands + NuGet republishes: bump the Tensors ref and swap the local DpSgdStep mirror for the Tensors version

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features

    • Added optional GPU-accelerated fused training across contrastive learning, forecasting, graph, physics, VAE, GAN, and synthetic-data workflows, with automatic fallback.
    • Added batched diffusion noise application and prediction with independent timesteps per sample.
    • Added reusable differential privacy training support with gradient clipping and noise aggregation.
    • Added support for training additional learnable tensors alongside model layers.
  • Bug Fixes

    • Improved gradient-penalty training so discriminator updates receive the intended gradients.
    • Prevented duplicate parameter updates in Fourier neural operator training.

Kicks off the non-time-series GPU-residency sweep. Adds a shared helper for
classes that don't inherit from NeuralNetworkBase or TimeSeriesModelBase but
still want their Train() to route forward + backward + optimizer through the
compiled fused plan (weights / activations / Adam moments resident on-device
across the whole step).

* src/Training/GpuResidentFusedStep.cs — shared helper. TryResolveOptimizerConfig
  maps a runtime IGradientBasedOptimizer to the fused-plan OptimizerType +
  hyperparameters (Adam / AdamW / SGD supported via case-insensitive class-name
  match + reflection over Options.InitialLearningRate / Beta1 / Beta2 / Epsilon
  / WeightDecay). TryStep is the one-shot entry.

Wired the following single-net models through the helper (mirrors
NeuralNetworkBase.TrainWithFusedStep):

* GraphClassificationModel (NN base + GNN + pooling + cross-entropy) —
  routes tape training through the fused plan when float + DirectGpu +
  compilation are live; falls through to the existing eager tape+optimizer
  loop on any failure.
* LinkPredictionModel (NN base + GNN + node embeddings + BCE) — same
  wire-up as GraphClassificationModel.
* FourierNeuralOperator (lift → FourierLayers → project, hardcoded SGD +
  learning rate 0.001) — captures the whole spectral-conv chain in a fused
  SGD plan; falls through to the in-place SGD loop below.

GAN-family generators (CTGAN, DPCTGAN, PATEGAN, ...), CLAPModel (learned
scalar outside layer hierarchy), AutoDiffTabGenerator (per-sample variable
shape defeats plan replay) and other non-conforming shapes are queued for
follow-up commits. The shared helper is the common seam so they all wire
the same way once their per-step structure fits the fused-plan contract
(constant shape + Adam/AdamW/SGD optimizer + ITrainableLayer-carried params).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings July 10, 2026 12:44
@vercel

vercel Bot commented Jul 10, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
aidotnet_website Ready Ready Preview, Comment Jul 10, 2026 11:15pm
aidotnet-playground-api Ready Ready Preview, Comment Jul 10, 2026 11:15pm

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@coderabbitai

coderabbitai Bot commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 23 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: efc7e0ea-ca0b-4481-8a00-076ceded2c08

📥 Commits

Reviewing files that changed from the base of the PR and between 1e9aa26 and 8bc5a47.

📒 Files selected for processing (35)
  • Directory.Packages.props
  • src/Audio/Fingerprinting/CLAPModel.cs
  • src/Diffusion/DiffusionModelBase.cs
  • src/Diffusion/NoisePredictors/NoisePredictorBase.cs
  • src/Diffusion/Schedulers/NoiseSchedulerBase.cs
  • src/Finance/Forecasting/Foundation/CSDI.cs
  • src/Finance/Forecasting/Foundation/TFC.cs
  • src/Finance/Forecasting/Foundation/TOTEM.cs
  • src/Finance/Forecasting/Neural/MQCNN.cs
  • src/NeuralNetworks/NeuralNetworkBase.cs
  • src/NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs
  • src/NeuralNetworks/SyntheticData/CTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/FinDiffGenerator.cs
  • src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs
  • src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs
  • src/NeuralNetworks/SyntheticData/MisGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/TVAEGenerator.cs
  • src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs
  • src/NeuralNetworks/SyntheticData/TabSynGenerator.cs
  • src/NeuralNetworks/SyntheticData/TableGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs
  • src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs
  • src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs
  • src/NeuralNetworks/WGANGP.cs
  • src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs
  • src/Training/CompiledTapeTrainingStep.cs
  • src/Training/DpSgdFusedStep.cs
  • src/Training/GpuResidentFusedStep.cs
  • src/Training/MultiSlotFusedStep.cs
  • src/Training/WganGpFusedStep.cs

Walkthrough

The PR adds GPU-resident fused training paths with eager fallbacks across multiple model families, introduces shared DP-SGD processing and WGAN-GP higher-order gradient support, and updates diffusion training for per-element batched timesteps, noise application, and prediction.

Changes

GPU-resident training

Layer / File(s) Summary
Fused-step infrastructure
src/Training/GpuResidentFusedStep.cs, src/Training/CompiledTapeTrainingStep.cs
Adds GPU availability and optimizer resolution, fused execution routing, optional extra tensors, and deduplicated parameter collection.
General model training integration
src/Audio/Fingerprinting/CLAPModel.cs, src/Finance/Forecasting/..., src/NeuralNetworks/Tasks/Graph/*, src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs
Adds fused forward/loss/optimizer paths for CLAP, forecasting, graph, and Fourier neural operator training with tape-based fallbacks.
Synthetic-model fused training
src/NeuralNetworks/SyntheticData/*Generator.cs
Adds fused generator, discriminator, VAE, GOGGLE, TabSyn, and TimeGAN training branches while retaining fallback paths.
Gradient penalty and DP-SGD integration
src/NeuralNetworks/SyntheticData/{CTGAN,CausalGAN,CopulaGAN,DPCTGAN,TableGAN}Generator.cs, src/Training/DpSgdStep.cs, src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs
Records gradient-penalty graphs for outer differentiation and centralizes DP-SGD per-example gradient clipping, aggregation, and noise injection.

Batched diffusion noise handling

Layer / File(s) Summary
Batched diffusion contracts
src/Diffusion/Schedulers/NoiseSchedulerBase.cs, src/Diffusion/NoisePredictors/NoisePredictorBase.cs, src/Diffusion/DiffusionModelBase.cs
Adds validated batched noise addition and noise prediction helpers with optional per-element conditioning.
Per-element diffusion training
src/Diffusion/DiffusionModelBase.cs
Samples independent timesteps for batched inputs, applies batched noise, routes batched prediction and recomputation accordingly, and preserves rank-one behavior.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ModelTrain
  participant GpuResidentFusedStep
  participant CompiledTapeTrainingStep
  participant Optimizer
  ModelTrain->>GpuResidentFusedStep: Submit trainable layers, tensors, forward, loss, optimizer
  GpuResidentFusedStep->>CompiledTapeTrainingStep: Compile fused forward/backward step
  CompiledTapeTrainingStep->>Optimizer: Apply fused parameter update
  Optimizer-->>CompiledTapeTrainingStep: Return loss
  CompiledTapeTrainingStep-->>ModelTrain: Success or fallback status
  ModelTrain->>ModelTrain: Run GradientTape fallback when fused execution fails
Loading
sequenceDiagram
  participant DiffusionTrain
  participant NoiseSchedulerBase
  participant NoisePredictorBase
  participant TapeStepContext
  DiffusionTrain->>DiffusionTrain: Sample timestep and noise per batch element
  DiffusionTrain->>NoiseSchedulerBase: AddNoiseBatched
  NoiseSchedulerBase-->>DiffusionTrain: Batched noisy input
  DiffusionTrain->>NoisePredictorBase: PredictNoiseBatched
  NoisePredictorBase-->>DiffusionTrain: Batched predicted noise
  DiffusionTrain->>TapeStepContext: Recompute batched forward
Loading

Poem

GPUs hum softly, tapes wait in line,
Fused steps carry gradients fine.
Diffusion drifts with timesteps bright,
GAN critics find their missing light.
Batches now move as one—
Training’s faster dance has begun.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Out of Scope Changes check ⚠️ Warning The PR adds many fused-training, diffusion, and DP-SGD changes unrelated to the linked createGraph fix, so it exceeds the issue scope. Split the unrelated fused-training, diffusion, and DP-SGD work into separate PRs, and keep this one limited to the five gradient-penalty fixes.
Docstring Coverage ⚠️ Warning Docstring coverage is 77.55% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and matches the main theme, though it understates the broader set of fused-training changes.
Linked Issues check ✅ Passed All five WGAN-GP discriminator updates now use createGraph: true in the inner gradient tape, so the penalty can backpropagate into discriminator weights.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/gpu-resident-nonts

Comment @coderabbitai help to get the list of available commands.

franklinic and others added 2 commits July 10, 2026 09:02
…sters

Wires TFC, TOTEM, CSDI, MQCNN through the compiled fused SGD plan (all four use
in-place SGD with lr=0.001 as their eager path). Extends the shared helper with
an inline IsGpuResidentAvailable check so callers don't need to duplicate the
availability gate.

* GpuResidentFusedStep — adds IsGpuResidentAvailable static property (mirrors
  TimeSeriesModelBase.CanTrainOnGpu / NeuralNetworkBase.CanTrainOnGpu). TryStep
  short-circuits when unavailable so callers stay on their eager path.
* TFC — supervised + contrastive branches fused into one closure; the fused
  plan captures both losses in a single backward.
* TOTEM — reconstruction + VQ commitment terms fused; the recompute-loss
  closure re-derives the commitment loss on each replay for graph parity.
* CSDI — denoising-score-matching: the forward closure re-samples timestep +
  noise each step (via ComputeDenoisingPairTape) so replay produces fresh
  training data even though the plan's captured graph shape is fixed.
* MQCNN — multi-quantile pinball loss captured directly via ComputeMultiQuantilePinballLossTape.

All four fall through to the existing in-place SGD path on any fused-plan failure
(unsupported op, non-compilable graph, etc). Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…hetic generators + TVAE

Wires the generator side of the dual-network training loop through the fused
compiled plan for 9 SyntheticData generators. Each covers a distinct pattern:

* PATEGAN — per-sample loop; fused plan compiles on first sample, replays
  across the batch with fresh noise per iteration.
* MedSynth — batched noise -> decoder -> discriminator-frozen -> non-saturating
  log-sigmoid loss.
* CausalGAN — batched noise -> generator -> optional causal-structure ->
  output activations -> disc-frozen -> -avgFake.
* OCTGAN — per-sample loop; noise -> gen -> disc-frozen embedding -> SVDD dist².
* TableGAN — noise + real batch -> gen -> composite loss (fake-scores +
  information-loss + optional classification).
* CTGAN — the tricky one: pack (noise, cond, mask) into one persistent input
  tensor so the closure can slice cond and mask back out on replay; captures
  both -avgFake AND conditional cross-entropy on the fused plan.
* DPCTGAN — same pattern as CTGAN but simpler (no mask, no CE).
* CopulaGAN — same pattern as DPCTGAN.
* TVAE — encoder + reparam + decoder + composite ELBO (recon + KL) all fused;
  Reparameterize re-samples inside the closure so training stays stochastic.

Discriminator/critic/student layers are NOT passed to the fused step (their
weights stay frozen on the gen step, matching the eager path semantics).
Falls through to the eager tape+optimizer path on any failure.

The discriminator STEPS of these GANs are not yet wired — they need a separate
pass because their loss involves gradient penalty (WGAN-GP) or per-example DP
clipping (DPCTGAN) which don't compose cleanly with the single-closure fused
step yet.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings July 10, 2026 13:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

…data gens

Two more SyntheticData generators covered — each covers a distinct pattern:

* TabSyn.TrainVAEBatch — per-row VAE ELBO training. Fused plan compiles on
  first row, replays across the batch with refreshed input per iteration.
  Reparameterize re-samples inside the closure so training stays stochastic.
* MisGAN.TrainDataGeneratorStep — dual-noise generator (data noise + mask
  noise) packed into a single persistent input tensor; closure slices them
  back out on replay. Loss = -E[D_x(fakeRow ⊙ fakeMask)].
* MisGAN.TrainMaskGeneratorStep — simpler single-noise version; per-sample
  loop with fused plan replay across the batch.

The DataDiscriminator step (WGAN-style critic + weight clipping) and mask
discriminator step aren't wired — they need separate plans and the
ClipWeights post-step interacts poorly with the fused optimizer's
in-place update. Follow-up.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings July 10, 2026 13:35
@vercel

vercel Bot commented Jul 10, 2026

Copy link
Copy Markdown

Deployment failed with the following error:

Resource is limited - try again in 24 hours (more than 100, code: "api-deployments-free-per-day").

Learn More: https://vercel.com/franklins-projects-02a0b5a0?upgradeToPro=build-rate-limit

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

…non-DP disc

* TimeGAN Phase 1 (TrainReconstructionStepBatched) — embedder + recovery
  reconstruction MSE. Straight single-network fused-step conversion.
* TimeGAN Phase 2 (TrainSupervisedStepBatched) — supervisor's next-step MSE
  in the embedded space, with the embedder frozen (its layers not in the
  trainable set for this step). Two-tensor closure (ht, htNext).
* MedSynth non-DP disc step — pack (real, fake) into a single input tensor
  along axis 0; slice back in the loss for BCE-real + BCE-fake. Generator
  runs OUTSIDE the fused plan so its weights stay frozen on the critic step.

TimeGAN Phase 3 (adversarial + supervised joint) and MedSynth DP-SGD paths
still use their eager tape — the WGAN-GP gradient penalty and per-example
DP clipping don't compose with the single-closure fused optimizer step.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings July 10, 2026 13:46

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

franklinic and others added 3 commits July 10, 2026 10:18
Extends the fused compiled-plan training step to accept raw trainable tensors
that aren't naturally carried by an ITrainableLayer (e.g. GOGGLE's soft
adjacency A; CLAP's learned _logTemperature scalar). Solves the SOLID
concerns raised in review of the ITrainableLayer-wrapper alternative:
* ISP — no forcing raw tensors to implement the full ILayer surface with
  no-op Forward / SetTrainingMode / GetParameterGradients members
* LSP — no risk of "layer" wrappers with identity Forward diverging from
  the layer contract's behavioral expectations elsewhere in the codebase

Wire-up:
* CompiledTapeTrainingStep.TryStepWithFusedOptimizer — new optional
  extraTensors param. Threaded into the dedup-aware parameter collection
  (CollectDeduplicatedParametersWithExtras) so extras get moment buffers,
  gradient accumulation, and GPU-residency in exactly the same code path
  that layer-carried params do.
* GpuResidentFusedStep.TryStep — plumbs extras through. A callsite with
  only extras (no layers) is a valid config for models whose whole
  trainable surface is raw tensors.

Dedup is by Tensor<T> reference across BOTH sources, so an extra tensor
that also happens to be layer-carried is registered exactly once — the
same shared/tied-weight protection the layer-only collector provides.

Enables Phase 4E (GOGGLE + CLAP wire-ups).
Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Consumes Phase 4A's extraTensors parameter to route models with raw
trainable state (outside the ITrainableLayer contract) through the fused
compiled plan.

* GOGGLE — soft adjacency A (initialised in InitializeModel, projected
  after each optimizer step via ProjectAdjacencyConstraints) is passed
  through extraTensors. The fused optimizer allocates its moment buffer,
  accumulates gradients, and applies the update in-place. Reparameterize
  re-samples inside the forward closure so training stays stochastic;
  ProjectAdjacencyConstraints runs after the fused step to keep the
  adjacency on the valid soft-adjacency manifold.
* CLAP — the learned _logTemperature scalar (Radford 2021 / Wu 2023
  contrastive-alignment temperature) goes through extraTensors. Both
  audio and text encoders are in the layers list; the fused step handles
  all three parameter classes uniformly.

Both fall through to the existing eager tape+optimizer path on any
failure (unsupported optimizer, non-compilable graph, etc).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ining

Adds industry-standard batched-per-element timestep API to the diffusion stack
(Ho et al. 2020 §3, HuggingFace diffusers reference). Previously
DiffusionModelBase.Train sampled ONE timestep for the whole batch — reducing
the signal-to-noise diversity of the training gradient. Canonical DDPM training
samples a distinct timestep per batch element.

* INoiseScheduler.AddNoiseBatched(cleanBatch, noiseBatch, timesteps) — default
  interface implementation delegates to the scalar AddNoise per element for
  backward compatibility. Concrete schedulers can override with a fused
  batched implementation.
* NoisePredictorBase.PredictNoiseBatched(noisyBatch, timesteps, conditioning) —
  virtual with default slice-then-call-scalar implementation. Subclasses that
  want a fused batched forward override this to keep the training loop on-device.
* DiffusionModelBase.PredictNoiseBatched — model-level counterpart with the
  same default slice-then-call-scalar behavior.
* DiffusionModelBase.Train detects batched vs rank-1 input and routes through
  the batched noise scheduler + batched predictor when input.Rank >= 2. Rank-1
  unbatched inputs stay on the scalar path for backward compatibility.

This is the foundation for the industry-exceeding "batched-per-element +
fused-resident" diffusion training path — see follow-up commits that wire
DiffusionModelBase.Train through the fused compiled plan.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 13

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/Diffusion/DiffusionModelBase.cs`:
- Around line 1077-1117: The isBatched check in the training flow incorrectly
treats every rank-2+ tensor as batch-first, breaking unbatched image and volume
inputs. Replace this rank-based inference with an explicit model/data batching
contract, such as a protected batch-axis hook or options setting, and default to
historical scalar behavior when no batch axis is declared. Use that contract
consistently for batchSize, timestep generation, noise creation,
AddNoiseBatched, and PredictNoiseBatched while preserving PredictNoise for
unbatched inputs.
- Around line 1081-1117: Update the optimizer re-evaluation closure in the
training method to use the full per-element timesteps vector: for batched inputs
call PredictNoiseBatched with timesteps, while preserving PredictNoise with the
scalar timestep for unbatched inputs. Ensure the TapeStepContext objective
matches the initial forward pass rather than always using timesteps[0].

In `@src/Diffusion/NoisePredictors/NoisePredictorBase.cs`:
- Around line 1154-1186: PredictNoiseBatched passes the full conditioning tensor
to each per-sample PredictNoise call. In PredictNoiseBatched, detect
conditioning whose leading dimension equals batchSize, create a per-sample
conditioning slice matching elem, and pass that slice for the current b; retain
the original conditioning unchanged when it is genuinely shared, and dispose any
temporary slices if required by Tensor ownership conventions.

In `@src/Finance/Forecasting/Foundation/CSDI.cs`:
- Around line 297-317: Remove the fused optimizer path from CSDI’s stochastic
denoising training, or refactor it so one timestep/noise pair is sampled outside
the compiled plan and passed consistently to both ForwardDenoise and
ComputeDenoiseLoss. Do not call ComputeDenoisingPairTape separately in those
callbacks. Also ensure optimizer selection and hyperparameters come from
_optimizer rather than hard-coded SGD values and 0.001f.

In `@src/Finance/Forecasting/Foundation/TFC.cs`:
- Around line 291-304: ComputeLossCombined currently uses the outer input for
ComputeContrastiveLossTape, causing stale traced data during replay; use the
persistent/traced input established by ForwardCombined and
CompiledTapeTrainingStep instead. Update the fused loss path so supervised and
contrastive terms both consume the current persistent batch, while preserving
the existing shape alignment and loss combination behavior.

In `@src/Finance/Forecasting/Foundation/TOTEM.cs`:
- Around line 344-354: ComputeLossCombined must reuse the commitment loss
produced by the initial training forward instead of calling
ForwardNativeForTrainingWithCommitment(input) again. Capture and retain that
commitment term when the first forward invokes VectorQuantize, then have
ComputeLossCombined add the cached value to the reconstruction loss, ensuring
the EMA codebook update occurs only once per step.

In `@src/Interfaces/INoiseScheduler.cs`:
- Around line 189-203: Validate that noiseBatch has the same rank and every
dimension as cleanBatch, not just the leading batch dimension, before
calculating perElement or accessing noiseSpan. Update the validation in the
relevant tensor noise-scheduling method near the existing batch-size checks,
throwing ArgumentException with nameof(noiseBatch) for any shape mismatch.
- Around line 187-218: Remove the default AddNoiseBatched implementation from
INoiseScheduler<T>, keeping the interface scalar-only for net471 compatibility.
Move the batching logic into an internal helper or extension method with an
explicit scheduler parameter, preserving its validation and per-batch behavior,
and update DiffusionModelBase to call that helper instead of the interface
member.

In `@src/NeuralNetworks/SyntheticData/CTGANGenerator.cs`:
- Around line 769-787: Eliminate the redundant generator forward in the Loss
closure by capturing the activated generator output produced by Fwd, including
the result of ApplyOutputActivationsBatched, in a local variable accessible to
Loss. Reuse that captured act when calling ConditionalCrossEntropy with the
sliced condition and mask tensors, and remove the repeated
GeneratorForwardWithResidualBatched and activation calls from Loss.

In `@src/NeuralNetworks/SyntheticData/TabSynGenerator.cs`:
- Around line 777-782: The Loss closure redundantly recomputes encoder outputs
already produced by Fwd. Capture the mean and logVar values from the Fwd
closure’s EncoderForwardOnTape/SplitEncoderOutput result, then have Loss reuse
those captured values when calling ComputeElboLossTape instead of running the
encoder again.
- Around line 783-795: The fused-training loop can silently skip rows when
TryStep returns false after fusedEngaged becomes true. Update the loop around
AiDotNet.Training.GpuResidentFusedStep<T>.TryStep so any failure after fusion
begins either processes the current and remaining rows through the eager path or
aborts the entire batch explicitly, rather than continuing and returning after
partial fusion; preserve correct handling for rows before fusion starts.

In `@src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs`:
- Around line 690-692: Remove the second foreach loop over _fourierLayers when
building allTrainable in the relevant method, since those layers are already
included in Layers; ensure each ITrainableLayer<T> is collected only once before
paramTensors is constructed.

In `@src/Training/GpuResidentFusedStep.cs`:
- Line 33: Change the declaration of GpuResidentFusedStep<T> from public to
internal, keeping its static nature and aligning its accessibility with
CompiledTapeTrainingStep<T> so it remains available only to in-assembly model
code.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 36fb055c-aef4-4178-b792-5acd02bb14f9

📥 Commits

Reviewing files that changed from the base of the PR and between b975b88 and cfd39ed.

📒 Files selected for processing (26)
  • src/Audio/Fingerprinting/CLAPModel.cs
  • src/Diffusion/DiffusionModelBase.cs
  • src/Diffusion/NoisePredictors/NoisePredictorBase.cs
  • src/Finance/Forecasting/Foundation/CSDI.cs
  • src/Finance/Forecasting/Foundation/TFC.cs
  • src/Finance/Forecasting/Foundation/TOTEM.cs
  • src/Finance/Forecasting/Neural/MQCNN.cs
  • src/Interfaces/INoiseScheduler.cs
  • src/NeuralNetworks/SyntheticData/CTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs
  • src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs
  • src/NeuralNetworks/SyntheticData/MisGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/TVAEGenerator.cs
  • src/NeuralNetworks/SyntheticData/TabSynGenerator.cs
  • src/NeuralNetworks/SyntheticData/TableGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs
  • src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs
  • src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs
  • src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs
  • src/Training/CompiledTapeTrainingStep.cs
  • src/Training/GpuResidentFusedStep.cs

Comment thread src/Diffusion/DiffusionModelBase.cs
Comment thread src/Diffusion/DiffusionModelBase.cs
Comment thread src/Diffusion/NoisePredictors/NoisePredictorBase.cs
Comment thread src/Finance/Forecasting/Foundation/CSDI.cs Outdated
Comment thread src/Finance/Forecasting/Foundation/TFC.cs
Comment thread src/NeuralNetworks/SyntheticData/CTGANGenerator.cs
Comment thread src/NeuralNetworks/SyntheticData/TabSynGenerator.cs
Comment thread src/NeuralNetworks/SyntheticData/TabSynGenerator.cs Outdated
Comment thread src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs Outdated
Comment thread src/Training/GpuResidentFusedStep.cs Outdated
…s (4F/4G/4H)

Consumes Tensors PR ooples/AiDotNet.Tensors#763 (compiled backward with
createGraph=true, DP-SGD helper, multi-slot persistent inputs).

## 4G — WGAN-GP correctness fix + fused disc steps (5 files)

Root cause fixed for AiDotNet issue #1844: every ComputeGradientPenalty
implementation used the inner GradientTape's ComputeGradients WITHOUT
createGraph=true. The inner backward's ops didn't record on the outer
tape, so inputGradients had no GradFn chain back to the discriminator
weights. Effect: the gradient penalty was included in the loss VALUE
but NOT in the disc's gradient — WGAN-GP silently degraded to plain
WGAN with no 1-Lipschitz enforcement.

Fix (5 files: CTGAN, DPCTGAN, CopulaGAN, CausalGAN, TableGAN): add
createGraph: true to the inner ComputeGradients so the outer tape can
differentiate the penalty through the disc weights.

Also wired 4 disc steps through the fused compiled plan
(CausalGAN, CTGAN, CopulaGAN, TableGAN — DPCTGAN's disc is DP so goes
via 4H). Packs (real, fake) into one persistent input along axis 0;
loss closure splits scores back out and computes wasserstein + λ·GP.
The fused path activates once Tensors PR #763 lands (its 4C change
removes the !createGraph gate at GradientTape.cs:727); until then
these wire-ups fall through to the (now-corrected) eager tape path.

## 4H — DP-SGD wire-ups (2 files)

Local AiDotNet.Training.DpSgdStep<T> — drop-in mirror of the same
helper in Tensors PR #763 (Engines.Training.DpSgdStep<T>). Enables
the wire-ups to land in this PR without waiting on the Tensors NuGet
publish. Both implementations enforce the Abadi 2016 §3 Algorithm 1
clip-BEFORE-aggregate order via their structure.

* DPCTGAN.TrainDiscriminatorStepBatchedDP — routes the per-example
  WGAN-GP + DP-SGD critic through DpSgdStep<T>. Objective: Wasserstein
  + λ·GP per example, clipped per-example, aggregated, noised, averaged.
* MedSynth.TrainDiscriminatorStepPerExampleDPSGD — routes the per-
  example non-saturating BCE + DP-SGD critic through the same helper.

Once Tensors PR #763 merges + NuGet publishes, the local mirror can
be swapped for the Tensors version by changing the `using` — the API
surface is identical by design.

## 4F — Batched-per-element diffusion foundation

DiffusionModelBase.Train samples per-batch-element timesteps for
rank ≥ 2 inputs (Ho et al. 2020 canonical pattern; HuggingFace
diffusers reference), routing through NoiseSchedulerBase.AddNoiseBatched
+ DiffusionModelBase.PredictNoiseBatched. Rank-1 unbatched inputs
keep the scalar path for backward compat.

net471 doesn't support default interface implementations, so
AddNoiseBatched lives on NoiseSchedulerBase<T> (as virtual with a
default per-element delegate) rather than on the INoiseScheduler<T>
interface. DiffusionModelBase's fused-training wire-up (via Tensors
PR #763's PersistentInputRegistry) is queued for the NuGet-bump
follow-up commit — the API foundation is already in place.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Addresses every unresolved review comment on the non-TS GPU-residency PR.

## Correctness fixes

* DiffusionModelBase.Train — RecomputeForward closure now dispatches
  isBatched ? PredictNoiseBatched(inp, timesteps) : PredictNoise(inp, timestep),
  so optimizer re-evaluation matches the recorded batched forward objective.
* NoisePredictorBase.PredictNoiseBatched — detect batch-aligned conditioning
  (leading dim == batchSize) and slice per element; preserve shared conditioning
  when the leading dim doesn't match. Fixes shape-mismatch / conditioning-shared
  bugs in classifier-free-guidance-style predictors.
* CSDI.Train — remove the fused-plan branch. ComputeDenoisingPairTape samples
  fresh (timestep, noise) each call, and calling it in both Fwd and Loss produces
  independent samples that don't match. The compiled plan can't refresh the RNG
  per replay either. Path stays on the eager tape until Tensors' PersistentInputRegistry
  (PR ooples/AiDotNet.Tensors#763) lands.
* TFC.Train — ComputeContrastiveLossTape now runs INSIDE ForwardCombined so it
  consumes the current-step persistent input (`inp`), not the outer `input` which
  would freeze at compile time. Closure-captured local; Loss reads it with a
  null-guard covering the Fwd-then-Loss ordering invariant.
* TOTEM.Train — capture the commitment tensor from ForwardNativeForTrainingWithCommitment's
  first call; reuse it in Loss instead of running the quantizer a second time.
  Without this, the compiled path performs an EMA SetCodebookValue update TWICE
  per step and diverges from the eager path.
* NoiseSchedulerBase.AddNoiseBatched — validate full shape parity (rank + every
  dim), not only the leading batch dim. Prevents indexing beyond noiseBatch's span
  when a caller passes [B, smaller...] shapes.

## Correctness cleanup

* CTGAN.TrainGeneratorStepBatched — capture (act, condFromInput, maskFromInput)
  from Fwd's single generator pass; reuse in Loss so the conditional-CE term
  doesn't re-run GeneratorForwardWithResidualBatched + ApplyOutputActivationsBatched
  (was doubling the per-step generator forward cost).
* TabSynGenerator.TrainVAEBatch — capture (mean, logVar) from Fwd's single encoder
  pass; reuse in Loss so the encoder doesn't run twice per row. Also: if the fused
  step fails mid-batch after fusedEngaged=true, drop to the eager path for the
  remaining rows (no row is silently skipped).
* FourierNeuralOperator.TapeTrainStep — remove the duplicate _fourierLayers loop
  in both the fused-path allTrainable collection AND the eager paramList. Fourier
  layers are already registered into Layers at construction, so iterating them
  again double-registered each Fourier parameter and drove the eager SGD to apply
  the update twice per step.

## API cleanup

* GpuResidentFusedStep<T> — public → internal. Aligns with CompiledTapeTrainingStep<T>
  which is already internal; keeps in-assembly training plumbing out of the public
  API surface.

Also fixed an inadvertent null-forgiving operator (`!`) usage that slipped into
the review-fix edits. All new nullable annotations use explicit null-guards with
descriptive InvalidOperationException on invariant violation (per CLAUDE.md's
"never use null-forgiving operators" rule).

INoiseScheduler default-interface issue (net471 incompatibility) was already
fixed in an earlier commit (AddNoiseBatched moved to NoiseSchedulerBase).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ooples pushed a commit that referenced this pull request Jul 10, 2026
Addresses every unresolved review comment on the non-TS GPU-residency PR.

## Correctness fixes

* DiffusionModelBase.Train — RecomputeForward closure now dispatches
  isBatched ? PredictNoiseBatched(inp, timesteps) : PredictNoise(inp, timestep),
  so optimizer re-evaluation matches the recorded batched forward objective.
* NoisePredictorBase.PredictNoiseBatched — detect batch-aligned conditioning
  (leading dim == batchSize) and slice per element; preserve shared conditioning
  when the leading dim doesn't match. Fixes shape-mismatch / conditioning-shared
  bugs in classifier-free-guidance-style predictors.
* CSDI.Train — remove the fused-plan branch. ComputeDenoisingPairTape samples
  fresh (timestep, noise) each call, and calling it in both Fwd and Loss produces
  independent samples that don't match. The compiled plan can't refresh the RNG
  per replay either. Path stays on the eager tape until Tensors' PersistentInputRegistry
  (PR ooples/AiDotNet.Tensors#763) lands.
* TFC.Train — ComputeContrastiveLossTape now runs INSIDE ForwardCombined so it
  consumes the current-step persistent input (`inp`), not the outer `input` which
  would freeze at compile time. Closure-captured local; Loss reads it with a
  null-guard covering the Fwd-then-Loss ordering invariant.
* TOTEM.Train — capture the commitment tensor from ForwardNativeForTrainingWithCommitment's
  first call; reuse it in Loss instead of running the quantizer a second time.
  Without this, the compiled path performs an EMA SetCodebookValue update TWICE
  per step and diverges from the eager path.
* NoiseSchedulerBase.AddNoiseBatched — validate full shape parity (rank + every
  dim), not only the leading batch dim. Prevents indexing beyond noiseBatch's span
  when a caller passes [B, smaller...] shapes.

## Correctness cleanup

* CTGAN.TrainGeneratorStepBatched — capture (act, condFromInput, maskFromInput)
  from Fwd's single generator pass; reuse in Loss so the conditional-CE term
  doesn't re-run GeneratorForwardWithResidualBatched + ApplyOutputActivationsBatched
  (was doubling the per-step generator forward cost).
* TabSynGenerator.TrainVAEBatch — capture (mean, logVar) from Fwd's single encoder
  pass; reuse in Loss so the encoder doesn't run twice per row. Also: if the fused
  step fails mid-batch after fusedEngaged=true, drop to the eager path for the
  remaining rows (no row is silently skipped).
* FourierNeuralOperator.TapeTrainStep — remove the duplicate _fourierLayers loop
  in both the fused-path allTrainable collection AND the eager paramList. Fourier
  layers are already registered into Layers at construction, so iterating them
  again double-registered each Fourier parameter and drove the eager SGD to apply
  the update twice per step.

## API cleanup

* GpuResidentFusedStep<T> — public → internal. Aligns with CompiledTapeTrainingStep<T>
  which is already internal; keeps in-assembly training plumbing out of the public
  API surface.

Also fixed an inadvertent null-forgiving operator (`!`) usage that slipped into
the review-fix edits. All new nullable annotations use explicit null-guards with
descriptive InvalidOperationException on invariant violation (per CLAUDE.md's
"never use null-forgiving operators" rule).

INoiseScheduler default-interface issue (net471 incompatibility) was already
fixed in an earlier commit (AddNoiseBatched moved to NoiseSchedulerBase).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (5)
src/Diffusion/DiffusionModelBase.cs (2)

812-838: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Blocking: same missing-validation gap as NoisePredictorBase.PredictNoiseBatched.

noisyBatch/timesteps are unchecked for null, noisyBatch.Rank == 0 throws an opaque negative-array-size error at new int[noisyBatch.Rank - 1], and batchSize == 0 throws an unlabeled DivideByZeroException at perElement = noisyBatch.Length / batchSize. NoiseSchedulerBase.AddNoiseBatched in the same cohort validates all of this up front — this method (and its NoisePredictorBase counterpart) should match that contract since it's a public virtual entry point.

As per path instructions, "missing validation of external inputs" is a blocking production-readiness item.

🛡️ Proposed fix
 public virtual Tensor<T> PredictNoiseBatched(Tensor<T> noisyBatch, int[] timesteps)
 {
+    if (noisyBatch is null) throw new ArgumentNullException(nameof(noisyBatch));
+    if (timesteps is null) throw new ArgumentNullException(nameof(timesteps));
+    if (noisyBatch.Rank == 0 || noisyBatch.Shape[0] <= 0)
+        throw new ArgumentException("noisyBatch must have a non-empty batch dimension.", nameof(noisyBatch));
     int batchSize = noisyBatch.Shape[0];
     if (timesteps.Length != batchSize)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/Diffusion/DiffusionModelBase.cs` around lines 812 - 838, Add upfront
validation to DiffusionModelBase.PredictNoiseBatched and its NoisePredictorBase
counterpart: reject null noisyBatch and timesteps with argument-specific
exceptions, reject rank-zero noisyBatch with a clear ArgumentException, and
reject zero batch size before division using the same contract and exception
style as NoiseSchedulerBase.AddNoiseBatched. Keep the existing timesteps-length
validation after these checks.

Source: Path instructions


1072-1144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff

Duplicated per-element slicing logic — consider consolidating the scheduler fallback.

The else branch here (manual per-element AddNoise loop for non-NoiseSchedulerBase schedulers) re-implements the exact same slicing/reassembly logic as NoiseSchedulerBase.AddNoiseBatched. Since the comment already notes the reason AddNoiseBatched can't live on INoiseScheduler<T> directly is net471's lack of default interface members, an extension method (public static Tensor<T> AddNoiseBatched<T>(this INoiseScheduler<T>, ...)) works fine on net471 and would let this call site do scheduler.AddNoiseBatched(...) unconditionally instead of duplicating the loop here.

Note also the previously-flagged rank-batching concern and the stale-timestep RecomputeForward issue both look resolved in this revision (rank contract addressed per commit cdeda3a; RecomputeForward at Line 1248 now branches on isBatched).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/Diffusion/DiffusionModelBase.cs` around lines 1072 - 1144, Consolidate
the duplicated fallback batching logic by adding a net471-compatible extension
method named AddNoiseBatched for INoiseScheduler<T>, implementing the
per-element slicing, scalar AddNoise calls, and tensor reassembly currently
duplicated in Train. Then update the isBatched branch in Train to call
_scheduler.AddNoiseBatched(input, noiseBatch, timesteps) unconditionally and
remove the NoiseSchedulerBase type check and manual fallback.
src/Diffusion/NoisePredictors/NoisePredictorBase.cs (1)

1158-1217: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Blocking: PredictNoiseBatched skips boundary validation that its sibling AddNoiseBatched performs.

noisyBatch and timesteps are never null-checked (noisyBatch.Shape[0] / timesteps.Length will NRE on null), noisyBatch.Rank == 0 throws inside new int[noisyBatch.Rank - 1] with an opaque negative-array-size error, and batchSize == 0 throws an unlabeled DivideByZeroException at perElement = noisyBatch.Length / batchSize. Compare with NoiseSchedulerBase.AddNoiseBatched (same PR), which validates null, rank, shape parity, and length before touching any data. This is a public virtual entry point users/subclasses can call directly, so it should fail with the same clear, actionable errors.

As per path instructions, "missing validation of external inputs" is explicitly called out as a blocking production-readiness issue.

🛡️ Proposed fix
 public virtual Tensor<T> PredictNoiseBatched(Tensor<T> noisyBatch, int[] timesteps, Tensor<T>? conditioning = null)
 {
+    if (noisyBatch is null) throw new ArgumentNullException(nameof(noisyBatch));
+    if (timesteps is null) throw new ArgumentNullException(nameof(timesteps));
+    if (noisyBatch.Rank == 0 || noisyBatch.Shape[0] <= 0)
+        throw new ArgumentException("noisyBatch must have a non-empty batch dimension.", nameof(noisyBatch));
     int batchSize = noisyBatch.Shape[0];
     if (timesteps.Length != batchSize)
         throw new System.ArgumentException(
             $"timesteps length {timesteps.Length} does not match batch size {batchSize}.",
             nameof(timesteps));
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/Diffusion/NoisePredictors/NoisePredictorBase.cs` around lines 1158 -
1217, Update PredictNoiseBatched to mirror AddNoiseBatched’s boundary validation
before accessing tensor data: reject null noisyBatch and timesteps with clear
argument errors, reject rank-0 noisyBatch, and reject zero batch size using
actionable parameter-specific messages. Also validate any required shape and
length invariants consistently with AddNoiseBatched, then perform per-element
calculations only after validation succeeds.

Source: Path instructions

src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs (2)

644-660: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Redundant full-batch forward + gradient-penalty computed only for loss reporting. Lines 646–660 run a full-batch DiscriminatorForwardBatched (×2) and ComputeGradientPenalty (which itself spins a nested tape + backward), but the result is used solely for lossValue at Line 689 — DpSgdStep.ComputeClippedAggregatedGradients re-derives all of this per example. That doubles the WGAN-GP forward/backward cost on every critic step. Mirror the MedSynth path and accumulate the scalar loss inside the perExampleLoss callback, dropping the eager block.

♻️ Suggested approach
// Remove the eager realScores/fakeScores/avgReal/avgFake/wasserstein/gp/lossTensor block.
// Accumulate loss inside the callback instead:
T lossSum = NumOps.Zero;
var noisedGrads = AiDotNet.Training.DpSgdStep<T>.ComputeClippedAggregatedGradients(
    batchSize: exampleCount,
    perExampleLoss: exIdx =>
    {
        // ... compute w + λ·gpEx as today ...
        var l = Engine.TensorAdd(w, Engine.TensorMultiplyScalar(gpEx, NumOps.FromDouble(_options.GradientPenaltyWeight)));
        if (l.Length > 0) lossSum = NumOps.Add(lossSum, l[0]);
        return l;
    },
    /* ... */);
T lossValue = NumOps.Divide(lossSum, NumOps.FromDouble(exampleCount));
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs` around lines 644 - 660,
Remove the eager full-batch discriminator and gradient-penalty calculations
around DiscriminatorForwardBatched and ComputeGradientPenalty. Following the
MedSynth pattern, compute each example’s Wasserstein and gradient-penalty loss
inside the perExampleLoss callback passed to
DpSgdStep<T>.ComputeClippedAggregatedGradients, accumulate its scalar into
lossSum, and derive lossValue by dividing by exampleCount after aggregation.

806-886: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Remove the dead DP-SGD helpers

src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs:597, 806, 1119 still contains the old parameter-noise path (ClipAndNoiseGradient, ComputePerExampleNoisedGradients, UpdateDiscriminatorParametersDP) even though TrainDiscriminatorStepBatchedDP now routes through AiDotNet.Training.DpSgdStep<T>.ComputeClippedAggregatedGradients. Keeping both implementations invites drift; delete the unused legacy path.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs` around lines 806 - 886,
Remove the obsolete DP-SGD parameter-noise implementation: delete
ClipAndNoiseGradient, ComputePerExampleNoisedGradients, and
UpdateDiscriminatorParametersDP from DPCTGANGenerator, along with any now-unused
helpers or references. Preserve TrainDiscriminatorStepBatchedDP’s usage of
AiDotNet.Training.DpSgdStep<T>.ComputeClippedAggregatedGradients and clean up
resulting unused usings or fields.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/Finance/Forecasting/Foundation/TFC.cs`:
- Around line 285-295: The fused TFC training path still freezes preprocessing
outputs during tracing because ApplyInstanceNormalization and
ComputeFrequencyRepresentation use detached host-side operations. Remove the
affected fused implementation around ForwardCombined and revert training to the
eager tape, or fully rewrite both preprocessing methods using traceable Engine
operations before enabling fusion; do not leave a path that silently reuses the
first batch’s values.

In `@src/Finance/Forecasting/Foundation/TOTEM.cs`:
- Around line 339-350: Remove the compiled fused-step implementation around
ForwardCombined and its captured commitment handling (the block spanning the
local ForwardCombined/related loss delegates), including any use of
ForwardNativeForTrainingWithCommitment in that path. Retain and use the eager
tape training path until VectorQuantize and its EMA/codebook updates are
implemented with traceable engine operations that execute exactly once per
batch.

---

Outside diff comments:
In `@src/Diffusion/DiffusionModelBase.cs`:
- Around line 812-838: Add upfront validation to
DiffusionModelBase.PredictNoiseBatched and its NoisePredictorBase counterpart:
reject null noisyBatch and timesteps with argument-specific exceptions, reject
rank-zero noisyBatch with a clear ArgumentException, and reject zero batch size
before division using the same contract and exception style as
NoiseSchedulerBase.AddNoiseBatched. Keep the existing timesteps-length
validation after these checks.
- Around line 1072-1144: Consolidate the duplicated fallback batching logic by
adding a net471-compatible extension method named AddNoiseBatched for
INoiseScheduler<T>, implementing the per-element slicing, scalar AddNoise calls,
and tensor reassembly currently duplicated in Train. Then update the isBatched
branch in Train to call _scheduler.AddNoiseBatched(input, noiseBatch, timesteps)
unconditionally and remove the NoiseSchedulerBase type check and manual
fallback.

In `@src/Diffusion/NoisePredictors/NoisePredictorBase.cs`:
- Around line 1158-1217: Update PredictNoiseBatched to mirror AddNoiseBatched’s
boundary validation before accessing tensor data: reject null noisyBatch and
timesteps with clear argument errors, reject rank-0 noisyBatch, and reject zero
batch size using actionable parameter-specific messages. Also validate any
required shape and length invariants consistently with AddNoiseBatched, then
perform per-element calculations only after validation succeeds.

In `@src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs`:
- Around line 644-660: Remove the eager full-batch discriminator and
gradient-penalty calculations around DiscriminatorForwardBatched and
ComputeGradientPenalty. Following the MedSynth pattern, compute each example’s
Wasserstein and gradient-penalty loss inside the perExampleLoss callback passed
to DpSgdStep<T>.ComputeClippedAggregatedGradients, accumulate its scalar into
lossSum, and derive lossValue by dividing by exampleCount after aggregation.
- Around line 806-886: Remove the obsolete DP-SGD parameter-noise
implementation: delete ClipAndNoiseGradient, ComputePerExampleNoisedGradients,
and UpdateDiscriminatorParametersDP from DPCTGANGenerator, along with any
now-unused helpers or references. Preserve TrainDiscriminatorStepBatchedDP’s
usage of AiDotNet.Training.DpSgdStep<T>.ComputeClippedAggregatedGradients and
clean up resulting unused usings or fields.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 4b12dd05-1f9e-4bfb-972f-48b8753b1077

📥 Commits

Reviewing files that changed from the base of the PR and between cfd39ed and 1e9aa26.

📒 Files selected for processing (16)
  • src/Diffusion/DiffusionModelBase.cs
  • src/Diffusion/NoisePredictors/NoisePredictorBase.cs
  • src/Diffusion/Schedulers/NoiseSchedulerBase.cs
  • src/Finance/Forecasting/Foundation/CSDI.cs
  • src/Finance/Forecasting/Foundation/TFC.cs
  • src/Finance/Forecasting/Foundation/TOTEM.cs
  • src/NeuralNetworks/SyntheticData/CTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs
  • src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs
  • src/NeuralNetworks/SyntheticData/TabSynGenerator.cs
  • src/NeuralNetworks/SyntheticData/TableGANGenerator.cs
  • src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs
  • src/Training/DpSgdStep.cs
  • src/Training/GpuResidentFusedStep.cs

Comment thread src/Finance/Forecasting/Foundation/TFC.cs Outdated
Comment thread src/Finance/Forecasting/Foundation/TOTEM.cs Outdated
… + consumer wire-ups

Adds three fused-plan training primitives (mirrors of AiDotNet.Tensors PR #763)
and wires the DP-SGD consumers through them:

* DpSgdFusedStep<T> — Abadi 2016 §3 Algorithm 1 per-example clip-before-aggregate.
  Runs each per-example forward+backward through a compiled plan (LR=0 so weights
  don't drift between replays), then computes the global L2 norm, clips, noises,
  and aggregates via vectorized IEngine ops (TensorMultiply/ReduceSum for the
  norm, TensorMultiplyScalar/TensorAdd for the accumulator, TensorRandomNormalInto
  for on-device Gaussian noise). Returns aggregated gradients so the caller's
  configured optimizer (Adam/AdamW/SGD/...) applies the update — no hardcoded
  SGD step inside.

* WganGpFusedStep<T> — Gulrajani 2017 WGAN-GP critic step. Composes
  E[D(fake)] − E[D(real)] + λ·(‖∇_x̃ D(x̃)‖₂ − 1)² inside one compiled plan with
  the inner ∇_x̃ D(x̃) recorded via createGraph=true so it differentiates into
  disc weights (issue #1844 fix). OnesLike uses vectorized Engine.TensorFill.

* MultiSlotFusedStep<T> — N-slot persistent input mechanism with plan cache
  keyed by composite shape + parameter identity. Slots refreshed via
  AsSpan().CopyTo(AsWritableSpan()) — no per-element loops on the hot path.

Consumer wire-ups (this PR):

* DPCTGANGenerator.TrainDiscriminatorStepBatchedDP — primary path routes through
  DpSgdFusedStep.TryStep, falls back to the existing ComputePerExampleNoisedGradients
  when the fused path can't engage (non-GPU host / compilation disabled).
* MedSynthGenerator.TrainDiscriminatorStepPerExampleDPSGD — same primary/fallback
  layering.
* Both legacy fallback loops (ComputePerExampleNoisedGradients + MedSynth eager)
  are now themselves vectorized: global-L2-norm via ReduceSum(g·g), clipped
  accumulation via TensorMultiplyScalar+TensorAdd, on-device Gaussian noise via
  TensorRandomNormalInto+TensorAdd — no scalar per-element loops anywhere on the
  DP-SGD path.

Codebase convention adopted:

* Class-scope `private static IEngine Engine => AiDotNetEngine.Current;` and
  `private static readonly INumericOperations<T> Ops = MathHelper.GetNumericOperations<T>();`
  on each fused-step class (matches the ~20 activation/optimizer/etc. bases).
  No IEngine or INumericOperations<T> threaded through method signatures.

Skipped (follow-ups filed):

* WGAN-GP consumer rewire → #1845 (blocked on IGradientBasedOptimizer<T> config
  accessor extension; consumers already correct via GpuResidentFusedStep + local
  createGraph=true ComputeGradientPenalty).
* Diffusion consumer rewire → #1846 (blocked on per-generator forward-refactor
  to move _timestepProjection inside compiled plan while treating raw timestep
  as a persistent slot).

Verification:

* net8.0 + net471 + net10.0 all build clean on both AiDotNet and AiDotNet.Tensors.
* Old DpSgdStep.cs deleted (replaced by DpSgdFusedStep).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Commit messages auto-fixed

One or more commit messages did not follow Conventional Commits, so they were rewritten to comply (subject case, header length ≤ 100, valid type). Each commit and its diff were preserved — no squashing.

The branch was force-pushed with the corrected messages. If you have local work on this branch, run git pull --rebase (or reset to the remote) before pushing again.

@ooples ooples force-pushed the feat/gpu-resident-nonts branch from e68b0d6 to a799a2d Compare July 10, 2026 20:16
ooples added a commit to ooples/AiDotNet.Tensors that referenced this pull request Jul 10, 2026
…ersistent inputs (#763)

* feat(training): compiled backward + createGraph, DP-SGD helper, multi-slot persistent inputs

Three engine primitives that unblock AiDotNet's Phase 4 GPU-residency work
(PR ooples/AiDotNet#1843) for training patterns the compiled fused plan
couldn't previously handle.

## 4C — compiled-backward path supports createGraph=true

GradientTape.ComputeGradients previously gated the compiled-backward path
on !createGraph (line 727), so higher-order AD workloads — Hessian-vector
products, WGAN-GP gradient penalty, meta-learning inner loops — always fell
back to the slow tape-walking backward even on persistent tapes.

Fix: allow the compiled path when createGraph=true, keeping the outer tape
active (don't SetCurrentTape(null)) and setting _isBackwardCreateGraph so
BackwardFunctions' AccumulateGrad uses out-of-place TensorAdd (keeps the
double-backward graph intact).

Regression test: CompiledBackwardCreateGraphTests validates first-order
gradient equivalence (createGraph=true vs false produces identical outputs)
and second-order AD flow (the WGAN-GP-style pattern — grad-of-grad wrt a
weight tensor — produces a non-zero, finite result). Before the fix, the
second-order test would silently return all zeros because inputGradients
from the inner tape had no GradFn chain back to the disc weights.

## 4D — DP-SGD helper (Abadi 2016 Algorithm 1)

New Engines.Training.DpSgdStep<T>.ComputeClippedAggregatedGradients wraps
the canonical differentially-private SGD step:
  - per-example forward+backward on an inner tape
  - GLOBAL L2 norm across ALL concatenated parameter gradients per example
  - per-example clip by min(1, C / ||g_i||_2) — BEFORE aggregation, which is
    the L2-sensitivity bound the DP proof requires (reversing the order
    breaks the privacy guarantee)
  - aggregate + Gaussian noise ~ N(0, σ² C² I) + average by batch size

The clip-then-aggregate order is enforced by the helper's structure — callers
can't accidentally aggregate-then-clip. Enables MedSynth DP disc and
DPCTGAN.TrainDiscriminatorStepBatchedDP in the consumer PR to route through
a shared, tested primitive instead of each re-implementing the same recipe.

Regression tests: three properties validated —
  - high clip / no noise reproduces plain-mean SGD (isolates the aggregation
    contract from the noise-injection contract)
  - tight clip forces each per-example gradient to L2 = clipNorm
  - asymmetric gradients across parameters expose the difference between
    per-example clip and aggregate-then-clip — this test would fail if the
    order were ever reversed

## 4B — multi-slot persistent-input registry

New Engines.Training.PersistentInputRegistry<T> generalizes the prior
two-slot (input + target) persistent-input mechanism to N slots. Enables
batched-per-element diffusion training (Ho et al. 2020, HuggingFace
diffusers reference) where each step provides three refreshed tensors:
noisy sample, target noise, and per-batch-element timesteps. Also useful
for classifier-free-guidance-style diffusion with multiple conditioning
streams, and TFT-style forecasters with multiple pooled inputs.

Contract: caller registers slots (each is a persistent Tensor<T> allocated
once), refreshes each slot's data per training step; the compiled plan's
captured reference stays stable so no recompile is needed across steps.

Regression tests: three properties validated —
  - reference stability + data-refresh (RefreshSlot preserves identity)
  - N slots of different shapes are independently refreshable
  - a compiled plan reads fresh data on Step() after RefreshSlot without
    invalidate-and-recompile (the batched-timestep end-to-end scenario)

Builds green net471 / net8.0 / net10.0. All 9 new tests pass. Zero regressions
in the existing compiled-training test suite.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): real fused primitives for multi-slot / DP-SGD / WGAN-GP disc

Supersedes the thin helpers (PersistentInputRegistry + eager DpSgdStep) in the
prior commit with real fused-plan-backed primitives that address the actual
"deferred work" concerns:

## MultiSlotFusedStep<T> (4B — real N-slot persistent inputs)

Not a wrapper over a list — a stateful trainer that owns its persistent input
tensors internally and drives a compiled plan against them. On each Step:
1. Composite-shape-key cache check triggers recompile if any slot's shape
   changed, or if the parameter set identity changed
2. Fresh caller data is copied into the persistent-reference tensors
3. The compiled plan captures those references at trace time; per-step
   replays re-read the refreshed data without recompilation
4. Parameters are marked GPU-resident on first Step (via .Gpu()) so the
   fused optimizer takes its on-device path

Enables the batched-per-element diffusion training pattern (Ho et al. 2020):
each step provides [clean_sample, noise, per-batch-element timesteps] as
three slots that all get refreshed together. Also unblocks classifier-free
guidance (latent + text_emb + class_emb + timesteps), TFT-style forecasters
with static/historical/future covariates, and CSDI's stochastic denoising.

## DpSgdFusedStep<T> (4D — real fused per-example DP-SGD)

Per-example forward+backward runs the COMPILED PLAN with LR=0 (so backward
populates .Grad without weights drifting between examples), then per-example
gradients are extracted, clipped against the GLOBAL L2 norm across all
parameters concatenated (Abadi 2016 §3 L2-sensitivity contract), summed,
Gaussian-noised, averaged, and applied via a final SGD-style update.

Fused benefit over the prior eager per-example loop: each example's
forward+backward now runs through the compiled plan (GPU-resident params,
fused kernels, no per-op host↔device round-trip) instead of a fresh
non-persistent GradientTape. The clip/aggregate/noise runs in host code
because its control flow (per-example L2 → min-clip against C) doesn't fit
the compiled-plan capture model — but these are O(params) scalar ops, not
the compute bottleneck. The forward+backward per example IS the bottleneck,
and it's fused.

The clip-BEFORE-aggregate order is enforced structurally so future edits
can't accidentally reverse it and silently break the DP privacy proof.

## WganGpFusedStep<T> (4C — fused WGAN-GP critic via nested-tape)

Full WGAN-GP critic objective — E[D(fake)] − E[D(real)] + λ·(‖∇_x̃ D(x̃)‖₂ − 1)²
— compiled into ONE fused plan. The inner ∇_x̃ D(x̃) gradient uses
createGraph=true so its ops record on the outer (compilation) tape;
combined with the createGraph=true compiled-backward support (prior commit),
the outer backward can differentiate the gradient penalty through the disc
weights. This is the fix for AiDotNet issue #1844 in its fused form.

Persistent slots: [real, fake, epsilon_01]. The trainer builds the
interpolated point x̃ = ε · real + (1-ε) · fake inside the traced graph so
per-step ε refreshes propagate into the fused plan without recompilation.
Parameters go GPU-resident on first Step.

Delivers the "compile the per-example loop inside the plan" (4D) and
"specialized WGAN-GP nested-tape primitive" (4C intermediate — real fused
disc without the general nested-tape compiler rewrite). Both address the
actual deferrals the previous helpers glossed over.

All 3 new primitives + all 10 existing Training-namespace tests pass on
net10.0. Builds clean net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): vectorize fused DP-SGD + WGAN-GP inner loops via IEngine ops

Replaces the per-element scalar loops in DpSgdFusedStep + WganGpFusedStep with
vectorized IEngine dispatches — the entire DP-SGD arithmetic (clip norm,
clipped accumulation, noise, aggregate) and the WGAN-GP scaffolding
(OnesLike, GP composition) now run through the same engine that owns the
compiled plan, so CPU SIMD / GPU kernels handle the tensor math instead of
host-side scalar iteration.

DpSgdFusedStep<T> (Abadi 2016):

* clippedSums accumulator zero-init: TensorFill(sums, Zero).
* Global L2 norm per parameter: sum(g²) via TensorMultiply(g,g) +
  ReduceSum(axes=null) — reduces across ALL axes to a scalar. Per-parameter
  scalar tensors accumulate (unavoidable since each parameter has its own
  gradient tensor and Abadi's L2-sensitivity contract requires the GLOBAL
  concatenated norm — cannot be a single tensor op).
* Clipped accumulation: TensorMultiplyScalar(g, clipFactor) + TensorAdd(sum, scaled).
* Noise + average: TensorRandomNormalInto for on-device Gaussian sampling,
  TensorMultiplyScalar(sum, 1/B) + TensorAdd(scaled, noise) for the DP-clean
  aggregate. Zero-noise path (noiseStd == 0) skips the noise tensor allocation.
* Returns `out Dictionary<Tensor<T>, Tensor<T>> aggregatedGradients` keyed by
  parameter tensor reference — caller applies with their configured
  optimizer (Adam/AdamW/SGD/...), preserving optimizer choice.
* `Random rng` parameter retained for API compatibility but noise dispatch
  now uses the engine's RNG (which supports vectorized / on-device draws).

WganGpFusedStep<T> (Gulrajani 2017):

* OnesLike helper switched from a scalar fill loop to Engine.TensorFill.
* Removed the redundant OnesLike-then-Subtract dance — now inline
  Engine.TensorSubtract(norm, OnesLike(norm)).

Codebase convention adopted:

* Class-scope `private static IEngine Engine => AiDotNetEngine.Current;` and
  `private static readonly INumericOperations<T> Ops = MathHelper.GetNumericOperations<T>();`
  on both classes (matches the ~20 activation/optimizer/etc. bases).
* No IEngine or INumericOperations<T> threaded through method signatures —
  BuildWganGpLoss and OnesLike now use the class-scope Engine / Ops directly.

Correctness contract unchanged:

* DP-SGD clip-BEFORE-aggregate order enforced by class structure.
* WGAN-GP createGraph=true on the inner tape so GP differentiates into
  disc weights (AiDotNet #1844 requirement).

Verification:

* net8.0, net471, net10.0 all build clean.
* API-compatible with the AiDotNet PR #1843 consumer wire-ups (DpSgdFusedStep
  used by DPCTGAN + MedSynth via a local mirror until the NuGet republishes).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

---------

Co-authored-by: franklinic <franklin@ivorycloud.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…e ops

Resolves the two BLOCKING CodeRabbit findings on PR #1843: both fused-plan
training paths were freezing preprocessing into the compiled plan because
their inner ops used host-side .Data.Span loops, so later replays reused the
first batch's normalized values / spectrum / argmin decisions instead of
recomputing per batch.

TFC — ApplyInstanceNormalization (RevIN) + ComputeFrequencyRepresentation:

* ApplyInstanceNormalization now delegates to a new stateless NormalizeWithStats
  that returns (normalized, mean, std) as tensors. Under the hood: ReduceMean +
  ReduceVariance + TensorSqrt + TensorBroadcastSubtract + TensorBroadcastDivide.
  All ops record on the tape and re-execute per replay under the compiled plan.
* DenormalizeForecast likewise delegates to DenormalizeForecastWithStats, which
  takes mean/std as explicit tensor parameters. ForwardNative threads them
  through as locals so the compile-mode replay uses CURRENT-step stats instead
  of frozen trace-time values.
* _revinMean/_revinStd (Vector<T> scalars) → _revinMeanTensor/_revinStdTensor
  (Tensor<T>? nullable) — kept for the abstract override's external callers,
  but the fused path never reads them.
* ComputeFrequencyRepresentation now uses Engine.RFFT for batched real FFT →
  reshape to [B, halfN, 2] pairs → TensorMultiply + ReduceSum(axis=2) for
  magSquared → TensorSqrt + TensorMultiplyScalar(1/n) for the one-sided
  magnitude spectrum → TensorSlice + TensorFlip + TensorConcatenate to mirror
  bins [1..n-halfN] into the tail. Handles even and odd n identically to the
  old scalar impl.
* Fused-plan fast path in TFC.Train restored — now safe because both
  preprocessing methods trace correctly.

TOTEM — VectorQuantize (VQ-VAE) traceable rewrite + post-Step EMA:

* New VectorQuantizeTraceable returns (quantized, commitmentLoss, argmin, head)
  using Engine ops end-to-end: TensorBroadcastSubtract + TensorMultiply +
  ReduceSum for distances, TensorArgMin along the codebookSize axis, per-c
  TensorSliceAxis + TensorIndexSelectDiff + TensorStack for the gather,
  Engine.StopGradient for the straight-through estimator, ReduceSum-based
  commitment loss weighted by β/totalLen.
* EMA moved OUT of the compiled forward into a new UpdateCodebookEMA(head,
  argmin). Called POST-Step by the fused path with the trace-time graph-node
  references — their .Data reflects the LAST replay so the update lands
  exactly once per batch (matches CodeRabbit's "EMA must execute exactly once
  per batch" contract). Under the compiled plan, argmin/head are refreshed
  by every _plan.Step() so post-Step reads see the current batch's values.
* UpdateCodebookEMA expresses the per-codebook scatter as: current codebook
  slice + TensorScatterAdd((1-decay)·(head - gathered), argmin) → new slice,
  then TensorConcatenate across codebook axis into the full [numCodebooks,
  codebookSize, codebookDim] tensor, then Engine.TensorCopy back into the
  _codebooks tensor object to preserve identity (future reads via the same
  reference see the update).
* Legacy VectorQuantize is now a thin adapter around VectorQuantizeTraceable
  for callers that don't need the extras.
* ForwardNativeForTrainingWithCommitment delegates to a new
  ForwardNativeForTrainingWithVQExtras that exposes the argmin/head; the
  original (forecast, commitmentLoss) contract is preserved for callers that
  don't need EMA state.
* Fused-plan fast path in TOTEM.Train restored — now safe because VQ is
  fully traceable AND the EMA runs exactly once per batch in post-Step eager
  code.

No new Tensors primitives were needed — the engine already has RFFT,
ReduceMean/Variance, TensorSqrt, TensorBroadcastSubtract/Divide, TensorFlip,
TensorConcatenate, TensorArgMin, TensorIndexSelectDiff, TensorSliceAxis,
TensorStack, StopGradient, TensorScatterAdd, and TensorCopy, all in the
autodiff/compile registry per OpRegistry.cs.

Verification:

* net8.0 + net471 + net10.0 all build clean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ooples pushed a commit that referenced this pull request Jul 10, 2026
…e ops

Resolves the two BLOCKING CodeRabbit findings on PR #1843: both fused-plan
training paths were freezing preprocessing into the compiled plan because
their inner ops used host-side .Data.Span loops, so later replays reused the
first batch's normalized values / spectrum / argmin decisions instead of
recomputing per batch.

TFC — ApplyInstanceNormalization (RevIN) + ComputeFrequencyRepresentation:

* ApplyInstanceNormalization now delegates to a new stateless NormalizeWithStats
  that returns (normalized, mean, std) as tensors. Under the hood: ReduceMean +
  ReduceVariance + TensorSqrt + TensorBroadcastSubtract + TensorBroadcastDivide.
  All ops record on the tape and re-execute per replay under the compiled plan.
* DenormalizeForecast likewise delegates to DenormalizeForecastWithStats, which
  takes mean/std as explicit tensor parameters. ForwardNative threads them
  through as locals so the compile-mode replay uses CURRENT-step stats instead
  of frozen trace-time values.
* _revinMean/_revinStd (Vector<T> scalars) → _revinMeanTensor/_revinStdTensor
  (Tensor<T>? nullable) — kept for the abstract override's external callers,
  but the fused path never reads them.
* ComputeFrequencyRepresentation now uses Engine.RFFT for batched real FFT →
  reshape to [B, halfN, 2] pairs → TensorMultiply + ReduceSum(axis=2) for
  magSquared → TensorSqrt + TensorMultiplyScalar(1/n) for the one-sided
  magnitude spectrum → TensorSlice + TensorFlip + TensorConcatenate to mirror
  bins [1..n-halfN] into the tail. Handles even and odd n identically to the
  old scalar impl.
* Fused-plan fast path in TFC.Train restored — now safe because both
  preprocessing methods trace correctly.

TOTEM — VectorQuantize (VQ-VAE) traceable rewrite + post-Step EMA:

* New VectorQuantizeTraceable returns (quantized, commitmentLoss, argmin, head)
  using Engine ops end-to-end: TensorBroadcastSubtract + TensorMultiply +
  ReduceSum for distances, TensorArgMin along the codebookSize axis, per-c
  TensorSliceAxis + TensorIndexSelectDiff + TensorStack for the gather,
  Engine.StopGradient for the straight-through estimator, ReduceSum-based
  commitment loss weighted by β/totalLen.
* EMA moved OUT of the compiled forward into a new UpdateCodebookEMA(head,
  argmin). Called POST-Step by the fused path with the trace-time graph-node
  references — their .Data reflects the LAST replay so the update lands
  exactly once per batch (matches CodeRabbit's "EMA must execute exactly once
  per batch" contract). Under the compiled plan, argmin/head are refreshed
  by every _plan.Step() so post-Step reads see the current batch's values.
* UpdateCodebookEMA expresses the per-codebook scatter as: current codebook
  slice + TensorScatterAdd((1-decay)·(head - gathered), argmin) → new slice,
  then TensorConcatenate across codebook axis into the full [numCodebooks,
  codebookSize, codebookDim] tensor, then Engine.TensorCopy back into the
  _codebooks tensor object to preserve identity (future reads via the same
  reference see the update).
* Legacy VectorQuantize is now a thin adapter around VectorQuantizeTraceable
  for callers that don't need the extras.
* ForwardNativeForTrainingWithCommitment delegates to a new
  ForwardNativeForTrainingWithVQExtras that exposes the argmin/head; the
  original (forecast, commitmentLoss) contract is preserved for callers that
  don't need EMA state.
* Fused-plan fast path in TOTEM.Train restored — now safe because VQ is
  fully traceable AND the EMA runs exactly once per batch in post-Step eager
  code.

No new Tensors primitives were needed — the engine already has RFFT,
ReduceMean/Variance, TensorSqrt, TensorBroadcastSubtract/Divide, TensorFlip,
TensorConcatenate, TensorArgMin, TensorIndexSelectDiff, TensorSliceAxis,
TensorStack, StopGradient, TensorScatterAdd, and TensorCopy, all in the
autodiff/compile registry per OpRegistry.cs.

Verification:

* net8.0 + net471 + net10.0 all build clean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Commit messages auto-fixed

One or more commit messages did not follow Conventional Commits, so they were rewritten to comply (subject case, header length ≤ 100, valid type). Each commit and its diff were preserved — no squashing.

The branch was force-pushed with the corrected messages. If you have local work on this branch, run git pull --rebase (or reset to the remote) before pushing again.

@ooples ooples force-pushed the feat/gpu-resident-nonts branch from 729ad3d to 87900aa Compare July 10, 2026 20:59
@vercel

vercel Bot commented Jul 10, 2026

Copy link
Copy Markdown

Deployment failed with the following error:

Resource is limited - try again in 1 day (more than 100, code: "api-deployments-free-per-day").

Learn More: https://vercel.com/franklins-projects-02a0b5a0?upgradeToPro=build-rate-limit

ooples pushed a commit that referenced this pull request Jul 10, 2026
…d primitives

Closes #1845 and #1846. Routes all four WGAN-GP critics through WganGpFusedStep
(the fused-plan primitive from PR #1843) and adds MultiSlotFusedStep wire-ups
to the three tabular diffusion consumers plus a base-class opt-in hook for
DiffusionModelBase subclasses.

## Optimizer config plumbing

Zero new API surface: IFusedOptimizerSpec.TryGetFusedOptimizerConfig already
exposes (OptimizerType, LR, Beta1, Beta2, Epsilon, WeightDecay, Schedule,
UseBf16Moments). Widened NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig
from private to internal so the sibling generators in the same assembly can
reuse the existing helper.

## WGAN-GP consumers (#1845)

* CTGANGenerator, CopulaGANGenerator, TableGANGenerator, CausalGANGenerator —
  each critic training method now attempts WganGpFusedStep.TryStep FIRST with
  the discriminator's optimizer hyperparameters extracted via
  TryMapToFusedOptimizerConfig, falls back to the existing
  GpuResidentFusedStep path (secondary fused), then the eager tape (final
  fallback). The ε ∈ [0, 1]^B epsilon sampler uses
  Engine.TensorRandomUniformRange to match each critic's local
  ComputeGradientPenalty behavior.
* Non-Adam optimizers (Lion, LBFGS) that don't implement IFusedOptimizerSpec
  cleanly fall through — TryMapToFusedOptimizerConfig returns false and the
  code path skips to GpuResidentFusedStep as before.

## Diffusion consumers (#1846)

* TabDDPMGenerator — refactored to expose a slot-based forward
  (BuildTabDDPMSlots + DenoiserForwardFromTensors +
  ComputeDiffusionLossTapeFromTensors). Per-row TrainBatch loop now attempts
  MultiSlotFusedStep with (numNoisy, actualNoise, catNoisy, catClean,
  rawSinusoidalTimeEmbed) as persistent slots. The learnable
  _timestepProjection stays INSIDE the compiled forward closure so its
  weights participate in the backward pass. Plan is compiled once on the
  first row and replayed via slot-data refresh for subsequent rows.
* TabSynGenerator — TrainDiffusionBatch's per-row loop wired with
  MultiSlotFusedStep on (noisyLatent, actualNoise, projectedTimeEmbed).
  Matches the existing eager path's semantic that _timestepProjection is
  NOT in _diffMLPLayers (kept detached in the eager path too), so the
  projected embedding is precomputed host-side per row and passed as slot
  data.
* Finance/Forecasting/Foundation/CSDI —
  - ApplyInstanceNormalization rewritten with traceable engine ops
    (ReduceMean + ReduceVariance + TensorSqrt + broadcast subtract/divide) —
    same pattern as the TFC RevIN fix. The previous `.Data.Span` per-batch
    loop froze at trace time.
  - New BuildCsdiSlots + DenoiserForwardFromSlots express the DDPM x_t
    formation and packed denoising input via TensorConcatenate + engine
    scalar multiplies. Replaces the `.Data.Span[i] = xt[0, i]` fill that
    baked the trace batch's x_t into the compiled plan.
  - Train() attempts MultiSlotFusedStep first, falls back to the existing
    eager ComputeDenoisingPairTape path when the fused path can't engage.
* Diffusion/DiffusionModelBase —
  - New opt-in `protected virtual bool SupportsFusedDenoising => false;`
    property. Base default is false so no existing subclass changes
    behavior.
  - Train() attempts MultiSlotFusedStep when SupportsFusedDenoising is true
    AND the training optimizer maps cleanly to a fused config. Slots:
    (noisySample, noise). Loss = MSE(pred, noise). QAT shadow restoration
    is preserved on the fused-success path.
  - Subclasses with fully-traceable PredictNoise / PredictNoiseBatched
    (e.g. after auditing to remove `.Data.Span` host loops) can opt in via
    a single-line override; no infrastructure changes needed elsewhere.

## Verification

* net8.0, net471, net10.0 all build clean.
* No API surface changes on IGradientBasedOptimizer<T> — the existing
  IFusedOptimizerSpec interface (already implemented by all fuse-able
  optimizers) provided everything needed.
* All consumers preserve eager fallback path for non-fuse-able optimizers
  and non-GPU hosts.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…d primitives (#1847)

Closes #1845 and #1846. Routes all four WGAN-GP critics through WganGpFusedStep
(the fused-plan primitive from PR #1843) and adds MultiSlotFusedStep wire-ups
to the three tabular diffusion consumers plus a base-class opt-in hook for
DiffusionModelBase subclasses.

## Optimizer config plumbing

Zero new API surface: IFusedOptimizerSpec.TryGetFusedOptimizerConfig already
exposes (OptimizerType, LR, Beta1, Beta2, Epsilon, WeightDecay, Schedule,
UseBf16Moments). Widened NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig
from private to internal so the sibling generators in the same assembly can
reuse the existing helper.

## WGAN-GP consumers (#1845)

* CTGANGenerator, CopulaGANGenerator, TableGANGenerator, CausalGANGenerator —
  each critic training method now attempts WganGpFusedStep.TryStep FIRST with
  the discriminator's optimizer hyperparameters extracted via
  TryMapToFusedOptimizerConfig, falls back to the existing
  GpuResidentFusedStep path (secondary fused), then the eager tape (final
  fallback). The ε ∈ [0, 1]^B epsilon sampler uses
  Engine.TensorRandomUniformRange to match each critic's local
  ComputeGradientPenalty behavior.
* Non-Adam optimizers (Lion, LBFGS) that don't implement IFusedOptimizerSpec
  cleanly fall through — TryMapToFusedOptimizerConfig returns false and the
  code path skips to GpuResidentFusedStep as before.

## Diffusion consumers (#1846)

* TabDDPMGenerator — refactored to expose a slot-based forward
  (BuildTabDDPMSlots + DenoiserForwardFromTensors +
  ComputeDiffusionLossTapeFromTensors). Per-row TrainBatch loop now attempts
  MultiSlotFusedStep with (numNoisy, actualNoise, catNoisy, catClean,
  rawSinusoidalTimeEmbed) as persistent slots. The learnable
  _timestepProjection stays INSIDE the compiled forward closure so its
  weights participate in the backward pass. Plan is compiled once on the
  first row and replayed via slot-data refresh for subsequent rows.
* TabSynGenerator — TrainDiffusionBatch's per-row loop wired with
  MultiSlotFusedStep on (noisyLatent, actualNoise, projectedTimeEmbed).
  Matches the existing eager path's semantic that _timestepProjection is
  NOT in _diffMLPLayers (kept detached in the eager path too), so the
  projected embedding is precomputed host-side per row and passed as slot
  data.
* Finance/Forecasting/Foundation/CSDI —
  - ApplyInstanceNormalization rewritten with traceable engine ops
    (ReduceMean + ReduceVariance + TensorSqrt + broadcast subtract/divide) —
    same pattern as the TFC RevIN fix. The previous `.Data.Span` per-batch
    loop froze at trace time.
  - New BuildCsdiSlots + DenoiserForwardFromSlots express the DDPM x_t
    formation and packed denoising input via TensorConcatenate + engine
    scalar multiplies. Replaces the `.Data.Span[i] = xt[0, i]` fill that
    baked the trace batch's x_t into the compiled plan.
  - Train() attempts MultiSlotFusedStep first, falls back to the existing
    eager ComputeDenoisingPairTape path when the fused path can't engage.
* Diffusion/DiffusionModelBase —
  - New opt-in `protected virtual bool SupportsFusedDenoising => false;`
    property. Base default is false so no existing subclass changes
    behavior.
  - Train() attempts MultiSlotFusedStep when SupportsFusedDenoising is true
    AND the training optimizer maps cleanly to a fused config. Slots:
    (noisySample, noise). Loss = MSE(pred, noise). QAT shadow restoration
    is preserved on the fused-success path.
  - Subclasses with fully-traceable PredictNoise / PredictNoiseBatched
    (e.g. after auditing to remove `.Data.Span` host loops) can opt in via
    a single-line override; no infrastructure changes needed elsewhere.

## Verification

* net8.0, net471, net10.0 all build clean.
* No API surface changes on IGradientBasedOptimizer<T> — the existing
  IFusedOptimizerSpec interface (already implemented by all fuse-able
  optimizers) provided everything needed.
* All consumers preserve eager fallback path for non-fuse-able optimizers
  and non-GPU hosts.

Co-authored-by: franklinic <franklin@ivorycloud.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
ooples and others added 2 commits July 10, 2026 19:55
… fused primitives (#1848)

Extends PR #1847 with three additional consumer wire-ups discovered during a
comprehensive audit against the primitives added in PR #1843. Each consumer
attempts the fused path first via the existing IFusedOptimizerSpec /
TryMapToFusedOptimizerConfig plumbing and falls back to eager tape training
when the fused path can't engage.

## WGAN-GP consumers (additional beyond #1845)

* NeuralNetworks/WGANGP.cs — standalone WGAN-GP class. Previously used the
  legacy flat-vector round-trip (three Critic.Predict calls per step,
  GetParameterGradients + host-side vector combine + UpdateCriticWithOptimizer).
  TrainCriticBatchWithGP now attempts WganGpFusedStep.TryStep first, using
  Critic.Layers as the parameter source and a per-batch ε ~ U(0, 1) sampler
  matching each critic's local ComputeGradientPenalty behavior. Falls back to
  the legacy path when the critic's optimizer has no fused-kernel mapping.

## Diffusion consumers (additional beyond #1846)

* NeuralNetworks/SyntheticData/FinDiffGenerator.cs — per-row DDPM training
  (Sattarov et al. 2023). TrainBatch caches a MultiSlotFusedStep across rows
  so the compiled plan is built once and replayed via slot-data refresh for
  subsequent rows. Slots: (packedDenoiserInput, targetNoise). Falls back to
  the existing Train(input, targetNoise) call on miss.

* NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs — per-row DDPM training
  with a custom TapeStepOver optimizer step. Same MultiSlotFusedStep caching
  pattern as FinDiff. Slots: (denoiserInput, targetNoise). Falls back to the
  existing tape-based TapeStepOver path on miss.

## Explicitly not wired in this PR (documented)

* NeuralNetworks/GenerativeAdversarialNetwork.cs — uses BCE loss with optional
  GP regularization (a separate auxiliary optimizer step), NOT Wasserstein +
  GP. Wiring WganGpFusedStep would change training semantics from BCE to
  Wasserstein. Requires a separate design decision.

* Finance/Forecasting/Foundation/{CCDM,MGTSD,TSDiff,TimeDiff,TimeGrad}.cs and
  Finance/Probabilistic/DiffusionTS.cs — each uses .Data.Span host-side packs
  for the denoising input (same trace-freeze issue TFC/CSDI hit in PR #1843).
  Each needs a TFC/CSDI-scale traceable rewrite BEFORE fused wiring is safe.
  Deferred to individual per-consumer PRs.

* MetaLearning/Algorithms/{MetaDDPMAlgorithm,MetaDMAlgorithm,MetaDiffAlgorithm}.cs
  — hand-rolled Vector&lt;T&gt; params with index arithmetic, NOT the
  ILayer/Engine/tape infrastructure. Would need full rewrite to fit MultiSlotFusedStep.

* DiffusionModelBase subclasses (DDPMModel, DiffWaveModel, LatentDiffusionModelBase)
  — the SupportsFusedDenoising opt-in hook (added in PR #1847) is available,
  but flipping it safely requires per-class audit of PredictNoise → UNet
  traceability. Left as follow-up work; the mechanism is in place.

## Verification

* net8.0, net471, net10.0 all build clean.
* No API surface changes.

Co-authored-by: franklinic <franklin@ivorycloud.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…for fused paths)

0.113.0 publishes Tensors PR #763 — this branch's gating dependency. #763's key fix
(4C): the compiled/persistent backward now honors createGraph=true (GradientTape.
ComputeGradients previously gated the compiled path on !createGraph), so the WGAN-GP
gradient penalty's inner backward differentiates into the disc weights through the
fused compiled plan instead of silently returning zeros (issue #1844). Against 0.111.2
the fused GPU-resident WGAN-GP path would have silently degraded to plain WGAN.

The src/Training fused-primitive mirrors (WganGpFusedStep, MultiSlotFusedStep,
DpSgdFusedStep) stay as the consumer-side primitives that #1847/#1848 centralize every
consumer through; they call the public engine API, so the bump alone routes them onto
#763's fixed compiled backward. Also picks up #765 (compiled ReduceMax axis fill) and
#764 (resident-param fused-Adam fix). Native OneDNN/OpenBLAS/CLBlast bumped in lockstep
(all published). Builds green on net8.0.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@ooples ooples force-pushed the feat/gpu-resident-nonts branch from b1a82df to f8fb2c9 Compare July 11, 2026 00:06
…onts

# Conflicts:
#	Directory.Packages.props
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

WGAN-GP gradient penalty not backpropagating into discriminator weights (createGraph=false in nested tape)

3 participants